"""
Type hint analysis for Python code.

Analyzes Python type hints and annotations to improve call graph accuracy
and provide better understanding of function signatures and data flow.
"""

import ast
from typing import Dict, List, Any, Optional, Set, Union, Tuple
from collections import defaultdict

from ...core.error_handler import AnalysisError, log_info


class TypeAnalyzer:
    """Analyzes Python type hints and annotations."""
    
    def __init__(self):
        """Initialize the type analyzer."""
        self.builtin_types = self._get_builtin_types()
        self.typing_module_types = self._get_typing_module_types()
        self.type_relationships = {}
    
    def _get_builtin_types(self) -> Set[str]:
        """Get set of built-in Python types."""
        return {
            'int', 'float', 'str', 'bool', 'bytes', 'list', 'dict', 'tuple',
            'set', 'frozenset', 'object', 'type', 'None', 'any', 'callable'
        }
    
    def _get_typing_module_types(self) -> Set[str]:
        """Get set of typing module types."""
        return {
            'Any', 'Union', 'Optional', 'List', 'Dict', 'Tuple', 'Set',
            'FrozenSet', 'Callable', 'Iterable', 'Iterator', 'Generator',
            'Coroutine', 'Awaitable', 'AsyncIterable', 'AsyncIterator',
            'Mapping', 'MutableMapping', 'Sequence', 'MutableSequence',
            'TypeVar', 'Generic', 'Protocol', 'Literal', 'Final',
            'ClassVar', 'NoReturn', 'Type', 'NamedTuple', 'TypedDict'
        }
    
    def analyze_types(self, ast_tree: ast.AST) -> Dict[str, Any]:
        """Analyze type hints and annotations in the AST."""
        type_info = {
            'type_annotations': [],
            'type_usage': defaultdict(int),
            'generic_types': [],
            'union_types': [],
            'optional_types': [],
            'callable_types': [],
            'custom_types': [],
            'type_relationships': {},
            'statistics': {
                'annotated_functions': 0,
                'annotated_variables': 0,
                'total_annotations': 0,
                'complex_types': 0,
                'generic_usage': 0
            }
        }
        
        visitor = TypeAnnotationVisitor(self.builtin_types, self.typing_module_types)
        visitor.visit(ast_tree)
        
        # Process collected annotations
        type_info['type_annotations'] = visitor.annotations
        type_info['generic_types'] = visitor.generic_types
        type_info['union_types'] = visitor.union_types
        type_info['optional_types'] = visitor.optional_types
        type_info['callable_types'] = visitor.callable_types
        type_info['custom_types'] = visitor.custom_types
        
        # Calculate usage statistics
        for annotation in visitor.annotations:
            type_name = annotation.get('type_name', '')
            type_info['type_usage'][type_name] += 1
            
            if annotation.get('annotation_type') == 'function':
                type_info['statistics']['annotated_functions'] += 1
            elif annotation.get('annotation_type') == 'variable':
                type_info['statistics']['annotated_variables'] += 1
        
        type_info['statistics']['total_annotations'] = len(visitor.annotations)
        type_info['statistics']['complex_types'] = len(visitor.generic_types) + len(visitor.union_types)
        type_info['statistics']['generic_usage'] = len(visitor.generic_types)
        
        # Analyze type relationships
        type_info['type_relationships'] = self._analyze_type_relationships(visitor.annotations)
        
        return type_info
    
    def _analyze_type_relationships(self, annotations: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Analyze relationships between types."""
        relationships = {
            'inheritance': [],
            'composition': [],
            'type_parameters': [],
            'return_type_patterns': defaultdict(list)
        }
        
        # Group annotations by function
        function_annotations = defaultdict(list)
        for annotation in annotations:
            if annotation.get('annotation_type') == 'function':
                func_name = annotation.get('function_name', '')
                function_annotations[func_name].append(annotation)
        
        # Analyze return type patterns
        for func_name, func_annotations in function_annotations.items():
            return_types = [ann for ann in func_annotations if ann.get('target_type') == 'return']
            arg_types = [ann for ann in func_annotations if ann.get('target_type') == 'argument']
            
            if return_types and arg_types:
                return_type = return_types[0].get('type_name', '')
                arg_type_names = [ann.get('type_name', '') for ann in arg_types]
                
                relationships['return_type_patterns'][return_type].append({
                    'function': func_name,
                    'argument_types': arg_type_names
                })
        
        return relationships
    
    def infer_types_from_usage(self, ast_tree: ast.AST) -> Dict[str, Any]:
        """Infer types from usage patterns where explicit annotations are missing."""
        inferred_types = {
            'variable_types': {},
            'function_return_types': {},
            'confidence_scores': {},
            'inference_methods': {}
        }
        
        visitor = TypeInferenceVisitor()
        visitor.visit(ast_tree)
        
        # Process inferred types
        for var_name, usage_patterns in visitor.variable_usage.items():
            inferred_type = self._infer_variable_type(usage_patterns)
            if inferred_type:
                inferred_types['variable_types'][var_name] = inferred_type
                inferred_types['confidence_scores'][var_name] = inferred_type.get('confidence', 0.0)
                inferred_types['inference_methods'][var_name] = inferred_type.get('method', 'unknown')
        
        return inferred_types
    
    def _infer_variable_type(self, usage_patterns: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
        """Infer variable type from usage patterns."""
        # Simple type inference based on usage patterns
        type_votes = defaultdict(int)
        confidence = 0.0
        
        for usage in usage_patterns:
            usage_type = usage.get('type')
            
            if usage_type == 'assignment':
                value_type = usage.get('value_type')
                if value_type:
                    type_votes[value_type] += 3  # High confidence for assignments
            
            elif usage_type == 'method_call':
                method = usage.get('method')
                if method in ['append', 'extend', 'pop', 'remove']:
                    type_votes['list'] += 2
                elif method in ['keys', 'values', 'items', 'get']:
                    type_votes['dict'] += 2
                elif method in ['add', 'remove', 'discard']:
                    type_votes['set'] += 2
            
            elif usage_type == 'iteration':
                type_votes['iterable'] += 1
            
            elif usage_type == 'subscript':
                type_votes['sequence'] += 1
        
        if type_votes:
            most_likely_type = max(type_votes, key=type_votes.get)
            confidence = type_votes[most_likely_type] / sum(type_votes.values())
            
            return {
                'type': most_likely_type,
                'confidence': confidence,
                'method': 'usage_analysis',
                'evidence': dict(type_votes)
            }
        
        return None
    
    def analyze_type_compatibility(self, source_type: str, target_type: str) -> Dict[str, Any]:
        """Analyze compatibility between two types."""
        compatibility = {
            'compatible': False,
            'requires_conversion': False,
            'conversion_method': None,
            'confidence': 0.0,
            'issues': []
        }
        
        # Normalize type names
        source_type = self._normalize_type_name(source_type)
        target_type = self._normalize_type_name(target_type)
        
        # Direct match
        if source_type == target_type:
            compatibility['compatible'] = True
            compatibility['confidence'] = 1.0
            return compatibility
        
        # Any type compatibility
        if target_type == 'Any' or source_type == 'Any':
            compatibility['compatible'] = True
            compatibility['confidence'] = 0.9
            return compatibility
        
        # Optional type handling
        if target_type.startswith('Optional['):
            inner_type = target_type[9:-1]  # Remove Optional[ and ]
            if source_type == inner_type or source_type == 'None':
                compatibility['compatible'] = True
                compatibility['confidence'] = 0.9
                return compatibility
        
        # Union type handling
        if target_type.startswith('Union['):
            union_types = self._parse_union_types(target_type)
            if source_type in union_types:
                compatibility['compatible'] = True
                compatibility['confidence'] = 0.9
                return compatibility
        
        # Built-in type conversions
        conversion_rules = {
            ('int', 'float'): {'compatible': True, 'conversion': 'implicit'},
            ('int', 'str'): {'compatible': True, 'conversion': 'str()'},
            ('float', 'str'): {'compatible': True, 'conversion': 'str()'},
            ('str', 'int'): {'compatible': True, 'conversion': 'int()', 'may_fail': True},
            ('str', 'float'): {'compatible': True, 'conversion': 'float()', 'may_fail': True},
            ('list', 'tuple'): {'compatible': True, 'conversion': 'tuple()'},
            ('tuple', 'list'): {'compatible': True, 'conversion': 'list()'},
        }
        
        rule_key = (source_type, target_type)
        if rule_key in conversion_rules:
            rule = conversion_rules[rule_key]
            compatibility['compatible'] = True
            compatibility['requires_conversion'] = rule.get('conversion') != 'implicit'
            compatibility['conversion_method'] = rule.get('conversion')
            compatibility['confidence'] = 0.7 if rule.get('may_fail') else 0.9
            
            if rule.get('may_fail'):
                compatibility['issues'].append('Conversion may fail at runtime')
        
        return compatibility
    
    def _normalize_type_name(self, type_name: str) -> str:
        """Normalize type name for comparison."""
        # Remove whitespace
        type_name = type_name.strip()
        
        # Handle typing module aliases
        aliases = {
            'typing.List': 'List',
            'typing.Dict': 'Dict',
            'typing.Tuple': 'Tuple',
            'typing.Set': 'Set',
            'typing.Optional': 'Optional',
            'typing.Union': 'Union',
            'typing.Any': 'Any',
            'typing.Callable': 'Callable'
        }
        
        return aliases.get(type_name, type_name)
    
    def _parse_union_types(self, union_type: str) -> List[str]:
        """Parse Union type to extract individual types."""
        # Simple parser for Union[Type1, Type2, ...]
        if union_type.startswith('Union[') and union_type.endswith(']'):
            inner = union_type[6:-1]  # Remove Union[ and ]
            # Split by comma, but be careful of nested brackets
            types = []
            current_type = ''
            bracket_depth = 0
            
            for char in inner:
                if char == ',' and bracket_depth == 0:
                    types.append(current_type.strip())
                    current_type = ''
                else:
                    current_type += char
                    if char in '[(':
                        bracket_depth += 1
                    elif char in ')]':
                        bracket_depth -= 1
            
            if current_type.strip():
                types.append(current_type.strip())
            
            return types
        
        return []


class TypeAnnotationVisitor(ast.NodeVisitor):
    """AST visitor for collecting type annotations."""
    
    def __init__(self, builtin_types: Set[str], typing_types: Set[str]):
        """Initialize the type annotation visitor."""
        self.builtin_types = builtin_types
        self.typing_types = typing_types
        self.annotations = []
        self.generic_types = []
        self.union_types = []
        self.optional_types = []
        self.callable_types = []
        self.custom_types = []
        self.current_function = None
        self.current_class = None
    
    def visit_ClassDef(self, node: ast.ClassDef):
        """Visit class definition."""
        old_class = self.current_class
        self.current_class = node.name
        self.generic_visit(node)
        self.current_class = old_class
    
    def visit_FunctionDef(self, node: ast.FunctionDef):
        """Visit function definition."""
        old_function = self.current_function
        self.current_function = node.name
        
        # Process function annotations
        self._process_function_annotations(node)
        
        self.generic_visit(node)
        self.current_function = old_function
    
    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
        """Visit async function definition."""
        old_function = self.current_function
        self.current_function = node.name
        
        self._process_function_annotations(node)
        
        self.generic_visit(node)
        self.current_function = old_function
    
    def visit_AnnAssign(self, node: ast.AnnAssign):
        """Visit annotated assignment."""
        if isinstance(node.target, ast.Name):
            var_name = node.target.id
            annotation = self._process_annotation(node.annotation)
            
            self.annotations.append({
                'annotation_type': 'variable',
                'variable_name': var_name,
                'target_type': 'variable',
                'type_name': annotation['type_name'],
                'type_category': annotation['category'],
                'raw_annotation': annotation['raw'],
                'line': node.lineno,
                'function': self.current_function,
                'class': self.current_class
            })
        
        self.generic_visit(node)
    
    def _process_function_annotations(self, node):
        """Process annotations for a function."""
        # Process argument annotations
        for arg in node.args.args:
            if arg.annotation:
                annotation = self._process_annotation(arg.annotation)
                
                self.annotations.append({
                    'annotation_type': 'function',
                    'function_name': node.name,
                    'target_type': 'argument',
                    'argument_name': arg.arg,
                    'type_name': annotation['type_name'],
                    'type_category': annotation['category'],
                    'raw_annotation': annotation['raw'],
                    'line': arg.lineno,
                    'function': self.current_function,
                    'class': self.current_class
                })
        
        # Process return annotation
        if node.returns:
            annotation = self._process_annotation(node.returns)
            
            self.annotations.append({
                'annotation_type': 'function',
                'function_name': node.name,
                'target_type': 'return',
                'type_name': annotation['type_name'],
                'type_category': annotation['category'],
                'raw_annotation': annotation['raw'],
                'line': node.lineno,
                'function': self.current_function,
                'class': self.current_class
            })
    
    def _process_annotation(self, annotation_node: ast.AST) -> Dict[str, Any]:
        """Process a single type annotation."""
        try:
            raw_annotation = ast.unparse(annotation_node)
        except:
            raw_annotation = '<unparseable>'
        
        type_info = {
            'raw': raw_annotation,
            'type_name': raw_annotation,
            'category': 'unknown'
        }
        
        # Analyze the annotation
        if isinstance(annotation_node, ast.Name):
            type_name = annotation_node.id
            type_info['type_name'] = type_name
            type_info['category'] = self._categorize_type(type_name)
        
        elif isinstance(annotation_node, ast.Attribute):
            # Handle qualified names like typing.List
            type_name = self._get_qualified_name(annotation_node)
            type_info['type_name'] = type_name
            type_info['category'] = self._categorize_type(type_name)
        
        elif isinstance(annotation_node, ast.Subscript):
            # Handle generic types like List[int], Optional[str]
            self._process_generic_annotation(annotation_node, type_info)
        
        elif isinstance(annotation_node, ast.Constant):
            # Handle string annotations (forward references)
            if isinstance(annotation_node.value, str):
                type_info['type_name'] = annotation_node.value
                type_info['category'] = 'string_annotation'
        
        return type_info
    
    def _process_generic_annotation(self, node: ast.Subscript, type_info: Dict[str, Any]):
        """Process generic type annotations."""
        base_type = self._get_annotation_name(node.value)
        type_info['type_name'] = type_info['raw']
        type_info['category'] = 'generic'
        
        # Categorize specific generic types
        if base_type == 'Optional':
            self.optional_types.append(type_info)
            type_info['category'] = 'optional'
        elif base_type == 'Union':
            self.union_types.append(type_info)
            type_info['category'] = 'union'
        elif base_type == 'Callable':
            self.callable_types.append(type_info)
            type_info['category'] = 'callable'
        elif base_type in self.typing_types:
            self.generic_types.append(type_info)
        else:
            self.custom_types.append(type_info)
            type_info['category'] = 'custom_generic'
    
    def _get_annotation_name(self, node: ast.AST) -> str:
        """Get the name of an annotation node."""
        if isinstance(node, ast.Name):
            return node.id
        elif isinstance(node, ast.Attribute):
            return node.attr
        else:
            return ''
    
    def _get_qualified_name(self, node: ast.Attribute) -> str:
        """Get qualified name from attribute node."""
        parts = []
        current = node
        
        while isinstance(current, ast.Attribute):
            parts.append(current.attr)
            current = current.value
        
        if isinstance(current, ast.Name):
            parts.append(current.id)
        
        return '.'.join(reversed(parts))
    
    def _categorize_type(self, type_name: str) -> str:
        """Categorize a type name."""
        if type_name in self.builtin_types:
            return 'builtin'
        elif type_name in self.typing_types:
            return 'typing'
        else:
            return 'custom'


class TypeInferenceVisitor(ast.NodeVisitor):
    """AST visitor for type inference from usage patterns."""
    
    def __init__(self):
        """Initialize the type inference visitor."""
        self.variable_usage = defaultdict(list)
        self.current_function = None
    
    def visit_FunctionDef(self, node: ast.FunctionDef):
        """Visit function definition."""
        old_function = self.current_function
        self.current_function = node.name
        self.generic_visit(node)
        self.current_function = old_function
    
    def visit_Assign(self, node: ast.Assign):
        """Visit assignment statement."""
        for target in node.targets:
            if isinstance(target, ast.Name):
                var_name = target.id
                value_type = self._infer_value_type(node.value)
                
                self.variable_usage[var_name].append({
                    'type': 'assignment',
                    'value_type': value_type,
                    'line': node.lineno,
                    'function': self.current_function
                })
        
        self.generic_visit(node)
    
    def visit_Call(self, node: ast.Call):
        """Visit function call."""
        # Check for method calls on variables
        if isinstance(node.func, ast.Attribute):
            if isinstance(node.func.value, ast.Name):
                var_name = node.func.value.id
                method_name = node.func.attr
                
                self.variable_usage[var_name].append({
                    'type': 'method_call',
                    'method': method_name,
                    'line': node.lineno,
                    'function': self.current_function
                })
        
        self.generic_visit(node)
    
    def visit_Subscript(self, node: ast.Subscript):
        """Visit subscript operation."""
        if isinstance(node.value, ast.Name):
            var_name = node.value.id
            
            self.variable_usage[var_name].append({
                'type': 'subscript',
                'line': node.lineno,
                'function': self.current_function
            })
        
        self.generic_visit(node)
    
    def visit_For(self, node: ast.For):
        """Visit for loop."""
        if isinstance(node.iter, ast.Name):
            var_name = node.iter.id
            
            self.variable_usage[var_name].append({
                'type': 'iteration',
                'line': node.lineno,
                'function': self.current_function
            })
        
        self.generic_visit(node)
    
    def _infer_value_type(self, value_node: ast.AST) -> Optional[str]:
        """Infer type from value node."""
        if isinstance(value_node, ast.Constant):
            return type(value_node.value).__name__
        elif isinstance(value_node, ast.List):
            return 'list'
        elif isinstance(value_node, ast.Dict):
            return 'dict'
        elif isinstance(value_node, ast.Tuple):
            return 'tuple'
        elif isinstance(value_node, ast.Set):
            return 'set'
        else:
            return None